import clickhouse_connect
from time import time

# 源ClickHouse配置（待迁移的cust_label库所在服务器）
source_config = {
    "host": "192.168.2.72",  # 例如：192.168.2.72
    "port": 8123,           # 通常是8123（HTTP端口）
    "user": "",       # 如无认证可留空
    "password": "",     # 如无认证可留空
    "database": "cust_label"
}

# 目标ClickHouse配置（192.168.3.5）
target_config = {
    "host": "192.168.3.35",
    "port": 8123,
    "user": "",
    "password": "",
    "database": "cust_label"  # 目标库名
}

def log(message):
    """带时间戳的日志输出"""
    print(f"[{time():.2f}] {message}")

def clean_create_sql(sql):
    """清理建表语句中的转义字符并规范换行/空白"""
    # 将 '\\n' '\\r' 转为实际换行，'\\t' 转空格
    sql = sql.replace('\\n', '\n').replace('\\r', '\n').replace('\\t', ' ')
    # 还原被转义的引号
    sql = sql.replace("\\'", "'").replace('\\"', '"')
    # 压缩多余空白，但保持换行边界
    import re
    sql = re.sub(r'[ \f\v]+', ' ', sql)
    sql = '\n'.join(line.strip() for line in sql.splitlines())
    return sql.strip()

def migrate_database():
    log("开始执行数据库迁移任务...")
    
    # 1. 连接源数据库
    try:
        source_client = clickhouse_connect.get_client(** source_config)
        log(f"成功连接源数据库：{source_config['host']}:{source_config['port']}/{source_config['database']}")
    except Exception as e:
        log(f"连接源数据库失败：{str(e)}")
        return

    # 2. 连接目标数据库并创建目标库（如果不存在）
    try:
        target_client = clickhouse_connect.get_client(**target_config)
        log(f"成功连接目标数据库：{target_config['host']}:{target_config['port']}")
        
        # 检查并创建目标数据库
        target_db = target_config["database"]
        db_exists = target_client.command(f"SELECT count() FROM system.databases WHERE name = '{target_db}'")
        if not int(db_exists):
            log(f"目标数据库 '{target_db}' 不存在，开始创建...")
            target_client.command(f"CREATE DATABASE {target_db}")
            log(f"目标数据库 '{target_db}' 创建成功")
        else:
            log(f"目标数据库 '{target_db}' 已存在，跳过创建")
    except Exception as e:
        log(f"连接目标数据库或创建数据库失败：{str(e)}")
        source_client.close()
        return

    try:
        # 3. 获取源库中所有表名
        log(f"开始获取源数据库 '{source_config['database']}' 中的表列表...")
        tables = source_client.command(
            "SELECT name FROM system.tables WHERE database = '{}'".format(source_config["database"]) 
        ).split('\n')
        tables = [t.strip() for t in tables if t.strip()]
        log(f"成功获取表列表，共 {len(tables)} 张表：{tables}")

        if not tables:
            log("源数据库中没有表，迁移结束")
            return

        # 4. 逐个迁移表
        total_start_time = time()
        for table_idx, table in enumerate(tables, 1):
            table_start_time = time()
            log(f"\n===== 开始迁移表 {table_idx}/{len(tables)}: {table} =====")

            # 4.1 获取并清理表结构
            try:
                create_sql = source_client.command(
                    f"SHOW CREATE TABLE {source_config['database']}.{table}"
                )
                # 清理非法字符
                cleaned_sql = clean_create_sql(create_sql)
                # 替换数据库名
                cleaned_sql = cleaned_sql.replace(
                    f"{source_config['database']}.{table}", 
                    f"{target_db}.{table}"
                )
                log(f"获取并清理表结构成功，处理后语句长度：{len(cleaned_sql)} 字符")
                log(f"处理后的建表语句前200字符：{cleaned_sql[:200]}...")  # 调试用
            except Exception as e:
                log(f"获取表 {table} 结构失败：{str(e)}，跳过该表")
                continue

            # 4.2 在目标库创建表（先删除已存在的表）
            try:
                log(f"开始在目标库创建表 {table}...")
                target_client.command(f"DROP TABLE IF EXISTS {target_config['database']}.{table}")
                # 使用清理后的DDL创建表
                target_client.command(cleaned_sql)
                log(f"表 {table} 结构创建成功")
            except Exception as e:
                log(f"创建表 {table} 结构失败：{str(e)}，跳过该表")
                continue

            # 4.3 读取源表数据并写入目标表（分批处理）
            try:
                # 获取可插入字段（排除 MATERIALIZED/ALIAS）
                col_rows = source_client.query(
                    f"SELECT name, default_kind FROM system.columns WHERE database = '{source_config['database']}' AND table = '{table}' ORDER BY position"
                ).result_rows
                columns = [name for name, default_kind in col_rows if default_kind not in ('MATERIALIZED', 'ALIAS')]
                columns_str = ', '.join([f"`{c}`" for c in columns])
                log(f"表 {table} 可插入字段 {len(columns)} 个：{', '.join(columns)}")

                # 分批迁移数据
                batch_size = 10000  # 每批处理1万条
                offset = 0
                total_rows = 0
                while True:
                    batch_start = time()
                    query = f"""
                        SELECT {columns_str} 
                        FROM {source_config['database']}.{table} 
                        LIMIT {batch_size} OFFSET {offset}
                    """
                    data = source_client.query(query).result_rows

                    if not data:
                        break  # 数据读取完毕

                    # 写入目标表
                    target_client.insert(
                        table=table,
                        database=target_config["database"],
                        data=data,
                        column_names=columns
                    )

                    # 统计进度
                    batch_rows = len(data)
                    total_rows += batch_rows
                    offset += batch_size
                    batch_time = time() - batch_start
                    log(f"已迁移 {total_rows} 条数据（本批：{batch_rows}条，耗时：{batch_time:.2f}秒）")

                table_time = time() - table_start_time
                log(f"===== 表 {table} 迁移完成，共 {total_rows} 条数据，耗时：{table_time:.2f}秒 =====")
            except Exception as e:
                log(f"迁移表 {table} 数据失败：{str(e)}，跳过该表")
                continue

        # 5. 迁移完成
        total_time = time() - total_start_time
        log(f"\n所有表迁移完成！总耗时：{total_time:.2f}秒")

    except Exception as e:
        log(f"迁移过程中发生错误：{str(e)}")
    finally:
        # 关闭连接
        source_client.close()
        target_client.close()
        log("数据库连接已关闭")

if __name__ == "__main__":
    migrate_database()